"""
ConnLLM价格管理
负责管理各个提供商和模型的价格信息，支持从多个位置加载
"""
import os
import json
import copy
import logging
from pathlib import Path
from typing import Dict, Any, Optional, Union, List, Tuple

# 设置日志
logger = logging.getLogger("connllm")

# 默认价格配置（美元/1000 tokens）
DEFAULT_PRICING = {
    "anthropic": {
        "claude-3-opus": {"input": 0.015, "output": 0.075},
        "claude-3-sonnet": {"input": 0.008, "output": 0.024},
        "claude-3-5-sonnet": {"input": 0.003, "output": 0.015},
        "claude-3-haiku": {"input": 0.00025, "output": 0.00125},
        "claude-3-7-sonnet": {"input": 0.015, "output": 0.075},
        "claude-3-7-sonnet:thinking": {"input": 0.015, "output": 0.075}
    },
    "openai": {
        "gpt-4o": {"input": 0.005, "output": 0.015},
        "gpt-4-turbo": {"input": 0.01, "output": 0.03},
        "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
        "o1": {"input": 0.005, "output": 0.015},
        "o1-mini": {"input": 0.0015, "output": 0.0045},
        "o1-preview": {"input": 0.01, "output": 0.03}
    },
    "openrouter": {
        "anthropic/claude-3-opus": {"input": 0.015, "output": 0.075},
        "anthropic/claude-3-sonnet": {"input": 0.008, "output": 0.024},
        "anthropic/claude-3-5-sonnet": {"input": 0.003, "output": 0.015},
        "anthropic/claude-3-haiku": {"input": 0.00025, "output": 0.00125},
        "anthropic/claude-3-7-sonnet": {"input": 0.015, "output": 0.075},
        "anthropic/claude-3-7-sonnet:thinking": {"input": 0.015, "output": 0.075},
        "openai/gpt-4o": {"input": 0.005, "output": 0.015},
        "openai/gpt-4-turbo": {"input": 0.01, "output": 0.03},
        "openai/gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
        "google/gemini-1.5-pro": {"input": 0.0005, "output": 0.0015},
        "google/gemini-1.5-flash": {"input": 0.00025, "output": 0.0005}
    },
    "moonshot": {
        "moonshot-v1-8k": {"input": 0.0001, "output": 0.0002},
        "moonshot-v1-32k": {"input": 0.0002, "output": 0.0004},
        "moonshot-v1-128k": {"input": 0.0004, "output": 0.0008},
        "moonshot-v1-auto": {"input": 0.0002, "output": 0.0004},
        "kimi-latest": {"input": 0.0002, "output": 0.0004}
    },
    "ppinfra": {
        "deepseek/deepseek-v3": {"input": 0.0001, "output": 0.0002},
        "deepseek/deepseek-r1": {"input": 0.0001, "output": 0.0002},
        "qwen/qwen-2.5-72b-instruct": {"input": 0.0001, "output": 0.0002}
    },
    "volces": {
        "deepseek-v3": {"input": 0.0001, "output": 0.0002},
        "deepseek-r1": {"input": 0.0001, "output": 0.0002},
        "doubao-1-5-pro-32k": {"input": 0.0001, "output": 0.0002}
    },
    "ollama": {
        # 本地模型，无需计费
    }
}


class PricingManager:
    """价格管理器，支持从多个位置加载价格信息"""
    
    def __init__(self):
        """初始化价格管理器"""
        self.pricing = copy.deepcopy(DEFAULT_PRICING)
        
        # 从多个位置加载价格信息
        self._load_pricing_from_files()
        
    def _load_pricing_from_files(self) -> None:
        """从多个位置加载价格信息文件"""
        # 价格文件搜索路径
        search_paths = self._get_pricing_search_paths()
        
        for path in search_paths:
            pricing_file = path / "connllm_pricing.json"
            if pricing_file.exists():
                try:
                    with open(pricing_file, "r", encoding="utf-8") as f:
                        loaded_pricing = json.load(f)
                        
                        # 验证格式
                        if not isinstance(loaded_pricing, dict):
                            logger.warning(f"价格文件格式错误: {pricing_file}")
                            continue
                            
                        # 更新价格
                        self._update_pricing(loaded_pricing)
                        
                    logger.debug(f"从 {pricing_file} 加载了价格信息")
                except (json.JSONDecodeError, IOError) as e:
                    logger.warning(f"无法加载价格文件 {pricing_file}: {str(e)}")
                    
        # 从旧的 pricing.json 加载 (兼容性)
        for path in search_paths:
            pricing_file = path / "pricing.json"
            if pricing_file.exists():
                try:
                    with open(pricing_file, "r", encoding="utf-8") as f:
                        loaded_pricing = json.load(f)
                        
                        # 验证格式
                        if not isinstance(loaded_pricing, dict):
                            logger.warning(f"旧格式价格文件格式错误: {pricing_file}")
                            continue
                            
                        # 更新价格
                        self._update_pricing(loaded_pricing)
                        
                    logger.debug(f"从旧格式文件 {pricing_file} 加载了价格信息")
                except (json.JSONDecodeError, IOError) as e:
                    logger.warning(f"无法加载旧格式价格文件 {pricing_file}: {str(e)}")
    
    def _get_pricing_search_paths(self) -> List[Path]:
        """
        获取价格文件搜索路径
        
        Returns:
            Path列表
        """
        paths = []
        
        # 当前工作目录
        paths.append(Path.cwd())
        
        # 用户配置目录
        if os.name == "nt":  # Windows
            app_data = os.environ.get("APPDATA")
            if app_data:
                paths.append(Path(app_data) / "ConnLLM")
        else:  # Linux/MacOS
            home = os.environ.get("HOME")
            if home:
                paths.append(Path(home) / ".config" / "connllm")
                
        # 项目目录 (如果作为包安装)
        try:
            import connllm
            package_dir = Path(connllm.__file__).parent
            paths.append(package_dir)
        except ImportError:
            pass
            
        return paths
        
    def _update_pricing(self, new_pricing: Dict[str, Dict[str, Dict[str, float]]]) -> None:
        """
        更新价格信息
        
        Args:
            new_pricing: 新的价格信息
        """
        for provider, models in new_pricing.items():
            if not isinstance(models, dict):
                logger.warning(f"提供商 {provider} 的价格格式错误")
                continue
                
            if provider not in self.pricing:
                self.pricing[provider] = {}
                
            for model, prices in models.items():
                if not isinstance(prices, dict):
                    logger.warning(f"模型 {provider}/{model} 的价格格式错误")
                    continue
                    
                # 确保至少有input和output字段
                if "input" not in prices or "output" not in prices:
                    logger.warning(f"模型 {provider}/{model} 的价格缺少input或output字段")
                    continue
                    
                self.pricing[provider][model] = prices
        
    def save_pricing(self) -> bool:
        """
        保存价格信息到用户配置目录
        
        Returns:
            是否成功保存
        """
        # 用户配置目录
        if os.name == "nt":  # Windows
            app_data = os.environ.get("APPDATA")
            if app_data:
                user_config_dir = Path(app_data) / "ConnLLM"
            else:
                return False
        else:  # Linux/MacOS
            home = os.environ.get("HOME")
            if home:
                user_config_dir = Path(home) / ".config" / "connllm"
            else:
                return False
                
        # 创建用户配置目录
        try:
            user_config_dir.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            logger.warning(f"无法创建用户配置目录: {str(e)}")
            return False
            
        # 价格文件路径
        pricing_file = user_config_dir / "connllm_pricing.json"
        
        try:
            # 保存价格信息
            with open(pricing_file, "w", encoding="utf-8") as f:
                json.dump(self.pricing, f, indent=4, ensure_ascii=False)
                
            logger.debug(f"价格信息已保存到 {pricing_file}")
            return True
        except Exception as e:
            logger.warning(f"无法保存价格信息到文件 {pricing_file}: {str(e)}")
            return False
            
    def get_price(self, provider: str, model: str) -> Dict[str, float]:
        """
        获取指定提供商和模型的价格
        
        Args:
            provider: 提供商名称
            model: 模型名称
            
        Returns:
            价格字典，包含input和output键
        """
        # 处理OpenRouter特殊情况
        if provider == "openrouter" and "/" in model:
            # OpenRouter模型格式为"provider/model"
            return self.pricing.get(provider, {}).get(model, {"input": 0.0, "output": 0.0})
            
        # 普通情况
        return self.pricing.get(provider, {}).get(model, {"input": 0.0, "output": 0.0})
        
    def calculate_cost(self, provider: str, model: str, input_tokens: int, output_tokens: int) -> float:
        """
        计算API调用成本
        
        Args:
            provider: 提供商名称
            model: 模型名称
            input_tokens: 输入令牌数
            output_tokens: 输出令牌数
            
        Returns:
            成本（美元）
        """
        prices = self.get_price(provider, model)
        input_cost = (input_tokens / 1000) * prices.get("input", 0.0)
        output_cost = (output_tokens / 1000) * prices.get("output", 0.0)
        return input_cost + output_cost
        
    def calculate_cost_from_response(self, provider: str, model: str, response: Dict[str, Any]) -> Tuple[float, int, int]:
        """
        从API响应中计算成本
        
        Args:
            provider: 提供商名称
            model: 模型名称
            response: API响应数据
            
        Returns:
            (成本（美元），输入令牌数，输出令牌数)元组
        """
        input_tokens = 0
        output_tokens = 0
        
        # 从不同提供商的响应中提取token使用量
        if provider == "anthropic":
            # Anthropic格式: response.message.usage.input_tokens/output_tokens
            if "message" in response and "usage" in response["message"]:
                usage = response["message"]["usage"]
                input_tokens = usage.get("input_tokens", 0)
                output_tokens = usage.get("output_tokens", 0)
        elif provider == "openai":
            # OpenAI格式: response.usage.prompt_tokens/completion_tokens
            if "usage" in response:
                usage = response["usage"]
                input_tokens = usage.get("prompt_tokens", 0)
                output_tokens = usage.get("completion_tokens", 0)
        elif provider == "openrouter":
            # OpenRouter格式: response.usage.prompt_tokens/completion_tokens
            if "usage" in response:
                usage = response["usage"]
                input_tokens = usage.get("prompt_tokens", 0)
                output_tokens = usage.get("completion_tokens", 0)
        # 其他提供商可以在此添加...
        
        cost = self.calculate_cost(provider, model, input_tokens, output_tokens)
        return cost, input_tokens, output_tokens
        
    def set_price(self, provider: str, model: str, input_price: float, output_price: float) -> None:
        """
        设置指定提供商和模型的价格
        
        Args:
            provider: 提供商名称
            model: 模型名称
            input_price: 输入价格（美元/1000 tokens）
            output_price: 输出价格（美元/1000 tokens）
        """
        if provider not in self.pricing:
            self.pricing[provider] = {}
            
        self.pricing[provider][model] = {
            "input": input_price,
            "output": output_price
        }
        
        # 保存到文件
        self.save_pricing()
        
    def get_provider_pricing(self, provider: str) -> Dict[str, Dict[str, float]]:
        """
        获取指定提供商的所有价格信息
        
        Args:
            provider: 提供商名称
            
        Returns:
            提供商的价格信息字典
        """
        return self.pricing.get(provider, {})
        
    def list_providers(self) -> List[str]:
        """
        获取所有提供商列表
        
        Returns:
            提供商名称列表
        """
        return list(self.pricing.keys())
        
    def list_models(self, provider: str) -> List[str]:
        """
        获取指定提供商的所有模型列表
        
        Args:
            provider: 提供商名称
            
        Returns:
            模型名称列表
        """
        return list(self.pricing.get(provider, {}).keys())
        
    def export_pricing(self, path: Union[str, Path]) -> bool:
        """
        导出价格信息到文件
        
        Args:
            path: 导出文件路径
            
        Returns:
            是否成功导出
        """
        try:
            with open(path, "w", encoding="utf-8") as f:
                json.dump(self.pricing, f, indent=4, ensure_ascii=False)
                
            logger.debug(f"价格信息已导出到 {path}")
            return True
        except Exception as e:
            logger.warning(f"无法导出价格信息到文件 {path}: {str(e)}")
            return False
            
    def import_pricing(self, path: Union[str, Path]) -> bool:
        """
        从文件导入价格信息
        
        Args:
            path: 导入文件路径
            
        Returns:
            是否成功导入
        """
        try:
            with open(path, "r", encoding="utf-8") as f:
                loaded_pricing = json.load(f)
                
            if not isinstance(loaded_pricing, dict):
                logger.warning(f"价格文件格式错误: {path}")
                return False
                
            # 更新价格
            self._update_pricing(loaded_pricing)
            
            # 保存到文件
            self.save_pricing()
            
            logger.debug(f"从 {path} 导入了价格信息")
            return True
        except Exception as e:
            logger.warning(f"无法从文件 {path} 导入价格信息: {str(e)}")
            return False
    
    def reset_pricing(self) -> None:
        """重置为默认价格"""
        self.pricing = copy.deepcopy(DEFAULT_PRICING)
        self.save_pricing()
        
    def get_formatted_price(self, provider: str, model: str) -> str:
        """
        获取格式化的价格字符串
        
        Args:
            provider: 提供商名称
            model: 模型名称
            
        Returns:
            格式化的价格字符串，如"$0.001/$0.002 (输入/输出 每千tokens)"
        """
        prices = self.get_price(provider, model)
        input_price = prices.get("input", 0.0)
        output_price = prices.get("output", 0.0)
        
        return f"${input_price:.6f}/${output_price:.6f} (输入/输出 每千tokens)"


# 全局价格管理器实例
_pricing_manager = PricingManager()

def get_pricing_manager() -> PricingManager:
    """
    获取全局价格管理器实例
    
    Returns:
        价格管理器实例
    """
    return _pricing_manager


def calculate_cost(provider: str, model: str, input_tokens: int, output_tokens: int) -> float:
    """
    计算API调用成本的快捷函数
    
    Args:
        provider: 提供商名称
        model: 模型名称
        input_tokens: 输入令牌数
        output_tokens: 输出令牌数
        
    Returns:
        成本（美元）
    """
    return get_pricing_manager().calculate_cost(provider, model, input_tokens, output_tokens)
